# coding=utf-8
# 函数式自动微分 https://www.mindspore.cn/tutorials/zh-CN/master/beginner/autograd.html
import numpy as np
import mindspore
from mindspore import nn
from mindspore import ops
from mindspore import Tensor, Parameter


x = ops.ones(5, mindspore.float32)  # input tensor
y = ops.zeros(3, mindspore.float32)  # expected output
w = Parameter(Tensor(np.random.randn(5, 3), mindspore.float32), name='w') # weight
b = Parameter(Tensor(np.random.randn(3,), mindspore.float32), name='b') # bias


def function(x, y, w, b):
    z = ops.matmul(x, w) + b
    loss = ops.binary_cross_entropy_with_logits(z, y, ops.ones_like(z), ops.ones_like(z))
    return loss

loss = function(x, y, w, b)
print(loss)

# 微分函数与梯度计算
grad_fn = mindspore.grad(function, (2, 3))
# 执行微分函数，即可获得w、b对应的梯度
grads = grad_fn(x, y, w, b)
print(grads)


# Stop Gradient
# 通常情况下，求导时会求loss对参数的导数，因此函数的输出只有loss一项。当我们希望函数输出多项时，微分函数会求所有输出项对参数的导数。
# 此时如果想实现对某个输出项的梯度截断，或消除某个Tensor对梯度的影响，需要用到Stop Gradient操作。
# 这里我们将function改为同时输出loss和z的function_with_logits，获得微分函数并执行。
def function_with_logits(x, y, w, b):
    z = ops.matmul(x, w) + b
    loss = ops.binary_cross_entropy_with_logits(z, y, ops.ones_like(z), ops.ones_like(z))
    return loss, z

grad_fn = mindspore.grad(function_with_logits, (2, 3))
grads = grad_fn(x, y, w, b)
print(grads)
# 此时如果想要屏蔽掉z对梯度的影响，即仍只求参数对loss的导数，可以使用ops.stop_gradient接口，将梯度在此处截断。我们将function实现加入stop_gradient，并执行。
def function_stop_gradient(x, y, w, b):
    z = ops.matmul(x, w) + b
    loss = ops.binary_cross_entropy_with_logits(z, y, ops.ones_like(z), ops.ones_like(z))
    return loss, ops.stop_gradient(z)
grad_fn = mindspore.grad(function_stop_gradient, (2, 3))
grads = grad_fn(x, y, w, b)
print(grads)

# Auxiliary data
# Auxiliary data意为辅助数据，是函数除第一个输出项外的其他输出。通常我们会将函数的loss设置为函数的第一个输出，其他的输出即为辅助数据。
# grad和value_and_grad提供has_aux参数，当其设置为True时，可以自动实现前文手动添加stop_gradient的功能，满足返回辅助数据的同时不影响梯度计算的效果。
# 下面仍使用function_with_logits，配置has_aux=True，并执行。
grad_fn = mindspore.grad(function_with_logits, (2, 3), has_aux=True)

grads, (z,) = grad_fn(x, y, w, b)
print(grads, z)


# 神经网络梯度计算
# Define model
class Network(nn.Cell):
    def __init__(self):
        super().__init__()
        self.w = w
        self.b = b

    def construct(self, x):
        z = ops.matmul(x, self.w) + self.b
        return z
# Instantiate model
model = Network()
# Instantiate loss function
loss_fn = nn.BCEWithLogitsLoss()

# Define forward function
def forward_fn(x, y):
    z = model(x)
    loss = loss_fn(z, y)
    return loss

# 由于使用Cell封装神经网络模型，模型参数为Cell的内部属性，此时我们不需要使用grad_position指定对函数输入求导，因此将其配置为None。
# 对模型参数求导时，我们使用weights参数，使用model.trainable_params()方法从Cell中取出可以求导的参数。
grad_fn = mindspore.value_and_grad(forward_fn, None, weights=model.trainable_params())
loss, grads = grad_fn(x, y)
print(grads)